import cv2
import time
import numpy as np
from random import randint
from rknn.api import RKNN



if __name__ == '__main__':

    rknn = RKNN()

    inWidth = 368
    inHeight = 256

    rknn.load_rknn('./final.rknn')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')

    cap = cv2.VideoCapture(10)

    hasFrame, frame = cap.read()
    fps = cap.get(cv2.CAP_PROP_FPS)

    # to decrease fps
    # normal read in 15 fps, 0.05s
    # computation time 0.2 s
    count = 0
    scale = 4

    while cv2.waitKey(1) < 0:
        t = time.time()


        hasFrame, frame = cap.read()
        count += 1
        if count is not scale:
            print("count = " + str(count))
            continue
        else:
            count = 0

        # q for end; space for stop
        key = cv2.waitKey(1) & 0xff
        if key == ord(" "):
            cv2.waitKey(0)
        if key == ord("q"):
            break

        # cv2.putText(frame, "FPS {0}".format(float('%.1f' % (counter / (time.time() - start_time)))), (500, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)

        # resize输入图像为368x368
        frame = cv2.resize(frame, (inWidth, inHeight), interpolation=cv2.INTER_CUBIC)
        if not hasFrame:
            cv2.waitKey()
            break
        frameWidth = frame.shape[1]
        frameHeight = frame.shape[0]

        # input mode转为’nchw’
        frame_input = np.transpose(frame, [2, 0, 1])
        t = time.time()
        [output] = rknn.inference(inputs=[frame_input], data_format="nchw")
        print("time:", time.time()-t)
    
        # rknn输出的数组转为1x57x46x46的矩阵
        output = output.reshape(1, 57, 46, 46)
    
        detected_keypoints = []
        keypoints_list = np.zeros((0,3))
        keypoint_id = 0
        threshold = 0.1

        for part in range(nPoints):
            probMap = output[0,part,:,:]
            probMap = cv2.resize(probMap, (frame.shape[1], frame.shape[0]))
            keypoints = getKeypoints(probMap, threshold)
            keypoints_with_id = []
            for i in range(len(keypoints)):
                keypoints_with_id.append(keypoints[i] + (keypoint_id,))
                keypoints_list = np.vstack([keypoints_list, keypoints[i]])
                keypoint_id += 1

            detected_keypoints.append(keypoints_with_id)


        frameClone = frame.copy()
    
        #for i in range(nPoints):
        #   for j in range(len(detected_keypoints[i])):
        #        cv2.circle(frameClone, detected_keypoints[i][j][0:2], 5, colors[i], -1, cv2.LINE_AA)
        #cv2.imshow("Keypoints",frameClone)
    
    
        valid_pairs, invalid_pairs = getValidPairs(output)
        personwiseKeypoints = getPersonwiseKeypoints(valid_pairs, invalid_pairs)
        #连接各个人体关键点
        for i in range(nPoints-1):
            for n in range(len(personwiseKeypoints)):
                index = personwiseKeypoints[n][np.array(POSE_PAIRS[i])]
                if -1 in index:
                    continue
                B = np.int32(keypoints_list[index.astype(int), 0])
                A = np.int32(keypoints_list[index.astype(int), 1])
                cv2.line(frameClone, (B[0], A[0]), (B[1], A[1]), colors[i], 3, cv2.LINE_AA)


        cv2.imshow("Detected Pose" , frameClone)



        
    
        #cv2.waitKey(0)

    rknn.release()